package com.wispx.gateway.core.filter.flow;

import cn.hutool.core.collection.ConcurrentHashSet;
import com.wispx.gateway.common.enums.FlowEnum;
import com.wispx.gateway.config.manager.DynamicConfigManager;
import com.wispx.gateway.config.manager.RouteListener;
import com.wispx.gateway.config.pojo.RouteDefinition;
import com.wispx.gateway.config.util.FilterConfigUtil;
import com.wispx.gateway.core.context.GatewayContext;
import com.wispx.gateway.core.filter.Filter;
import com.wispx.gateway.core.filter.flow.algorithm.LeakyBucketRateLimiter;
import com.wispx.gateway.core.filter.flow.algorithm.SlidingWindowRateLimiter;
import com.wispx.gateway.core.filter.flow.algorithm.TokenBucketRateLimter;

import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

import static com.wispx.gateway.common.constant.FilterConstant.FLOW_FILTER_NAME;
import static com.wispx.gateway.common.constant.FilterConstant.FLOW_FILTER_ORDER;

public class FlowFilter implements Filter {

    private Map<String/* 服务名 */, RateLimiter> rateLimiterMap = new ConcurrentHashMap<>();
    private Set<String/* 服务名 */> addedListenerSet = new ConcurrentHashSet<>();// 已添加监听器的服务集合

    @Override
    public void doPreFilter(GatewayContext ctx) {
        // 流控过滤器配置
        RouteDefinition.FlowFilterConfig flowFilterConfig = Optional.ofNullable(FilterConfigUtil.findFilterConfigByClass(ctx.getRouteDefinition().getFilterConfigs(), FLOW_FILTER_NAME, RouteDefinition.FlowFilterConfig.class))
                .orElse(new RouteDefinition.FlowFilterConfig());
        if (!flowFilterConfig.isEnable()) { // 没有开启流控
            ctx.doFilter();
        } else {
            String serviceName = ctx.getRouteDefinition().getServiceName();
            RateLimiter rateLimiter = rateLimiterMap.computeIfAbsent(serviceName, name -> {

                if (!addedListenerSet.contains(serviceName)) { // 没有监听过该服务
                    DynamicConfigManager.getInstance().addRouteListener(serviceName, new RouteListener() {
                        @Override
                        public void routeOnChange(RouteDefinition route) {
                            // 路由变化时，从rateLimiterMap移除该服务，以便重新获取RateLimiter
                            rateLimiterMap.remove(route.getServiceName());
                        }
                    });
                    addedListenerSet.add(serviceName);
                }

                return getRateLimiter(flowFilterConfig);
            });
            rateLimiter.tryConsume(ctx);
        }

    }

    @Override
    public void doPostFilter(GatewayContext ctx) {
        ctx.doFilter();
    }

    @Override
    public String mark() {
        return FLOW_FILTER_NAME;
    }

    @Override
    public int getOrder() {
        return FLOW_FILTER_ORDER;
    }

    @SuppressWarnings("DuplicateBranchesInSwitch")
    private RateLimiter getRateLimiter(RouteDefinition.FlowFilterConfig flowFilterConfig) {
        FlowEnum flowType = flowFilterConfig.getFlowType();
        switch (flowType) {
            case SLIDINGWINDOWS:
                return new SlidingWindowRateLimiter(flowFilterConfig.getCapacity(), flowFilterConfig.getRate());
            case LEAKYBUCKET:
                return new LeakyBucketRateLimiter(flowFilterConfig.getCapacity(), flowFilterConfig.getRate(), null);
            case TOKENBUCKET:
                return new TokenBucketRateLimter(flowFilterConfig.getCapacity(), flowFilterConfig.getRate());
            default:
                return new TokenBucketRateLimter(flowFilterConfig.getCapacity(), flowFilterConfig.getRate());
        }

    }
}
